{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyOmny2ODXVnj+MFw9vz8lFF",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Huangjian2013/ai-demo/blob/main/rag/05_Qdrant_02.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install qdrant-client --quiet\n",
        "!pip install sentence-transformers --quiet\n",
        "\n"
      ],
      "metadata": {
        "id": "moDSivjssz5Z"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_DNYVZBaOa9N"
      },
      "outputs": [],
      "source": [
        "from qdrant_client import QdrantClient\n",
        "from qdrant_client.models import PointStruct\n",
        "from qdrant_client.models import VectorParams, Distance\n",
        "import numpy\n",
        "from sentence_transformers import SentenceTransformer\n",
        "\n",
        "\n",
        "collection_name = \"embeded_collection\"\n",
        "encoder = SentenceTransformer(\"all-MiniLM-L6-v2\")\n",
        "\n",
        "client = QdrantClient(\":memory:\")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "docs = [\n",
        "    {\n",
        "        \"name\" : \"极越01\",\n",
        "        \"type\" : \"SUV\"\n",
        "    },\n",
        "    {\n",
        "        \"name\" : \"极越07\",\n",
        "        \"type\" : \"car\"\n",
        "    }\n",
        "    ]"
      ],
      "metadata": {
        "id": "-G67EsDsMwOc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "client.create_collection(\n",
        "  collection_name=collection_name,\n",
        "  vectors_config=VectorParams(\n",
        "      size=encoder.get_sentence_embedding_dimension(),\n",
        "      distance=Distance.COSINE),\n",
        ")"
      ],
      "metadata": {
        "collapsed": true,
        "id": "gF2RGmXEtBmA",
        "outputId": "27dd3180-935c-4da6-f000-61dbd82893a3",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 49
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "client.upload_collection(\n",
        "   collection_name=collection_name,\n",
        "   vectors=encoder.encode(docs),\n",
        "   payload=docs,\n",
        "   ids=None,  # Vector ids will be assigned automatically\n",
        "   batch_size=2,\n",
        "   parallel=1,\n",
        ")"
      ],
      "metadata": {
        "id": "j78Od4AmNdZt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "hits = client.search(\n",
        "    collection_name=collection_name,\n",
        "    query_vector=encoder.encode(\"01\"),\n",
        "    limit=1,\n",
        ")\n",
        "\n",
        "print(hits)"
      ],
      "metadata": {
        "id": "CTiNwFSLOcmq",
        "outputId": "8f6c16e4-228a-4f81-f5ee-592d098d4116",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[ScoredPoint(id='d076512c-8a3e-4fda-a315-cf7307df072c', version=0, score=0.6401567459106445, payload={'name': '极越01', 'type': 'SUV'}, vector=None, shard_key=None, order_value=None)]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "hits = client.search(\n",
        "    collection_name=collection_name,\n",
        "    query_vector=encoder.encode(\"07\"),\n",
        "    limit=1,\n",
        ")\n",
        "\n",
        "print(hits)"
      ],
      "metadata": {
        "id": "25woCytCN7nh",
        "outputId": "a07c427a-23bf-4387-b134-b90744c1e870",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[ScoredPoint(id='b6aa142b-ca4f-429b-814d-6c1c30fb5020', version=0, score=0.5518996715545654, payload={'name': '极越07', 'type': 'car'}, vector=None, shard_key=None, order_value=None)]\n"
          ]
        }
      ]
    }
  ]
}